#!/usr/bin/env python3
"""
异步下载TTS模型脚本
"""

import os
import sys
import asyncio
import logging
from concurrent.futures import ThreadPoolExecutor

# 设置环境变量
os.environ['XDG_CACHE_HOME'] = os.path.expanduser('~/.cache')
os.environ['XDG_DATA_HOME'] = os.path.expanduser('~/.local/share')
os.environ['TTS_CACHE_DIR'] = os.path.expanduser('~/.cache/tts')
os.environ['TORCH_HOME'] = os.path.expanduser('~/.cache/torch')
os.environ['HF_HOME'] = os.path.expanduser('~/.cache/huggingface')

# 创建目录
os.makedirs(os.path.expanduser('~/.cache/tts'), exist_ok=True)
os.makedirs(os.path.expanduser('~/.cache/torch'), exist_ok=True)
os.makedirs(os.path.expanduser('~/.cache/huggingface'), exist_ok=True)

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('model_download.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

try:
    from TTS.api import TTS
    print("TTS库导入成功")
except ImportError as e:
    print(f"TTS库导入失败: {e}")
    sys.exit(1)

# 优先级排序的模型列表
models_to_download = [
    # 高优先级 - 项目必需的模型
    ("tts_models/multilingual/multi-dataset/bark", "Bark - 多人对话必需", "high"),
    ("tts_models/zh-CN/baker/tacotron2-DDC-GST", "Baker - 中文专用模型", "high"),
    
    # 中优先级 - 项目支持的模型
    ("tts_models/multilingual/multi-dataset/xtts_v2", "XTTS v2 - 克隆语音", "medium"),
    
    # 低优先级 - 额外的模型
    ("tts_models/multilingual/multi-dataset/your_tts", "YourTTS - 高质量多语言", "low"),
    ("tts_models/en/ljspeech/vits", "VITS - 高质量英文模型", "low"),
    ("tts_models/en/ljspeech/tacotron2-DDC", "Tacotron2 - 稳定英文模型", "low"),
]

def download_single_model(model_info):
    """下载单个模型"""
    model_name, description, priority = model_info
    
    try:
        logger.info(f"🔄 开始下载 [{priority}]: {description}")
        
        # 检查模型是否已存在
        cache_dir = os.path.expanduser('~/.cache/tts')
        model_dir = os.path.join(cache_dir, model_name.replace('/', '--'))
        
        if os.path.exists(model_dir) and os.listdir(model_dir):
            logger.info(f"✅ 模型已存在，跳过: {description}")
            return True
        
        # 创建TTS实例，这会自动下载模型
        tts = TTS(model_name=model_name, progress_bar=False)
        
        logger.info(f"✅ 下载完成: {description}")
        return True
        
    except Exception as e:
        logger.error(f"❌ 下载失败: {description}")
        logger.error(f"错误信息: {str(e)[:200]}...")
        return False

async def download_models_by_priority():
    """按优先级下载模型"""
    
    # 按优先级分组
    high_priority = [m for m in models_to_download if m[2] == "high"]
    medium_priority = [m for m in models_to_download if m[2] == "medium"]
    low_priority = [m for m in models_to_download if m[2] == "low"]
    
    all_results = []
    
    # 使用线程池执行下载
    with ThreadPoolExecutor(max_workers=1) as executor:
        # 1. 先下载高优先级模型
        print("\n🔥 下载高优先级模型 (必需)...")
        high_tasks = [
            asyncio.get_event_loop().run_in_executor(executor, download_single_model, model)
            for model in high_priority
        ]
        high_results = await asyncio.gather(*high_tasks)
        all_results.extend(high_results)
        
        # 2. 再下载中优先级模型
        print("\n⚡ 下载中优先级模型...")
        medium_tasks = [
            asyncio.get_event_loop().run_in_executor(executor, download_single_model, model)
            for model in medium_priority
        ]
        medium_results = await asyncio.gather(*medium_tasks)
        all_results.extend(medium_results)
        
        # 3. 最后下载低优先级模型
        print("\n💫 下载低优先级模型...")
        low_tasks = [
            asyncio.get_event_loop().run_in_executor(executor, download_single_model, model)
            for model in low_priority
        ]
        low_results = await asyncio.gather(*low_tasks)
        all_results.extend(low_results)
    
    return all_results

async def main():
    """主函数"""
    print("="*60)
    print("🚀 TTS模型下载器 - 按优先级下载")
    print("="*60)
    
    try:
        results = await download_models_by_priority()
        
        # 统计结果
        success_count = sum(1 for r in results if r)
        total_count = len(results)
        
        print(f"\n✨ 下载完成统计:")
        print(f"✅ 成功: {success_count}/{total_count}")
        print(f"❌ 失败: {total_count - success_count}/{total_count}")
        
        # 显示下载位置
        cache_dir = os.path.expanduser('~/.cache/tts')
        print(f"\n📁 模型保存位置: {cache_dir}")
        
        # 列出已下载的模型
        if os.path.exists(cache_dir):
            print("\n📋 已下载的模型:")
            for item in sorted(os.listdir(cache_dir)):
                if os.path.isdir(os.path.join(cache_dir, item)):
                    print(f"  📂 {item}")
        
        if success_count >= 2:  # 至少成功下载2个核心模型
            print("\n🎉 核心模型下载成功，可以开始测试！")
        else:
            print("\n⚠️  核心模型下载不足，请检查网络连接")
            
    except KeyboardInterrupt:
        print("\n⏹️  用户中断下载")
    except Exception as e:
        print(f"\n💥 下载过程出错: {e}")

if __name__ == "__main__":
    asyncio.run(main())